// SPDX-License-Identifier: Apache-2.0
import { ConsoleLogger } from "@thi.ng/logger";
import {
	$xy,
	$y,
	F,
	V2,
	V4,
	defn,
	distance,
	div,
	float,
	ifThen,
	lte,
	madd,
	min,
	mix,
	mul,
	ret,
	smoothstep,
	step,
	sub,
	sym,
	ternary,
	vec2,
	type FloatSym,
	type TaggedFn1,
	type Vec2Sym,
	type Vec4Sym,
} from "@thi.ng/shader-ast";
import * as std from "@thi.ng/shader-ast-stdlib";
import { map, normRange2d, zip } from "@thi.ng/transducers";
import { glCanvas } from "@thi.ng/webgl";
import {
	shaderToy,
	type MainImageFn,
	type ShaderToyUniforms,
} from "@thi.ng/webgl-shadertoy";

// normalized margin for each function plot
const MARGIN = 0.18;
// grid layout config
const COLS = 4;
const ROWS = 4;

// list of easing functions to visualize
const EASINGS = [
	std.easeInBounce,
	std.easeOutBounce,
	std.easeInOutBounce,
	std.easeInOutElastic,
	std.easeInExpo,
	std.easeOutExpo,
	std.easeInOutExpo,
	std.easeInCubic,
	std.easeOutCubic,
	std.easeInOutCubic,
	std.easeInBack,
	std.easeOutBack,
	std.easeInOutBack,
	std.easeInQuart,
	std.easeOutQuart,
	std.easeInOutQuart,
];

// thi.ng/shader-ast typedefs for our custom shader uniforms
interface DemoUniforms extends ShaderToyUniforms {
	// curve thickness
	thickness: FloatSym;
	// dot radius for highlighting current position
	radius: FloatSym;
	// colors
	bgColor: Vec4Sym;
	bgHoverColor: Vec4Sym;
	curveColor1: Vec4Sym;
	curveColor2: Vec4Sym;
	dotColor: Vec4Sym;
}

// main shader function. the two args given are objects containing GLSL builtin
// vars and uniform bindings. IMPORTANT: all the terms & function calls used
// here are each only producing AST nodes, and the resulting syntax tree will
// later be compiled to GLSL. this approach opens up powerful meta-programming
// possibilities (using the full expressiveness of TypeScript) and indeed most
// of that main function body will be dynamically generated by transforming the
// above list of selected easing functions and pairing each with a small
// dedicated viewport rect to visualize those curves...
const main: MainImageFn<DemoUniforms> = (gl, unis) => {
	// predeclare local vars / symbols
	let uv: Vec2Sym, t: FloatSym;
	// the actual function body:
	return [
		// compute UV coordinate for current fragment (aka [0..1] range)
		// sym() is used to define a new symbol (variable) with the given inner expression as value
		// $xy() is a vector swizzle (e.g. in GLSL: gl_FragCoord.xy)
		(uv = sym(div($xy(gl.gl_FragCoord), unis.resolution))),
		// loop & slow down animation
		(t = sym(std.foldback01(mul(0.5, unis.time)))),
		// for each given easing function code we first generate a viewport
		// region check and an associated function call (incl. function
		// definition) to visualize that easing function...
		...map(
			([fn, pos]) =>
				// generate an if() statement
				ifThen(
					// conditional: check if we're in the correct region
					std.isPointInRect(
						uv,
						vec2(pos[0], pos[1]),
						vec2(1 / COLS, 1 / ROWS)
					),

					// the truthy branch of the conditional... if current
					// fragment is within the defined region, call the
					// visualization function and return its result.
					// META-PROGRAMMING ALERT: we first generate and then
					// immediately invoke the new shader function to visualize the
					// current easing function within the given region/rect.
					// this will result in just a function-call AST node with a
					// reference to the actual generated function being stored
					// as part of that node. during GLSL compilation,
					// thi.ng/shader-ast then produces a call graph and emits
					// all functions (and their transitive dependencies) in the
					// correct order...
					[
						ret(
							easingVisualization(
								fn,
								pos,
								[1 / COLS, 1 / ROWS],
								unis
							)(
								uv,
								$xy(gl.gl_FragCoord),
								div(unis.mouse, unis.resolution),
								t
							)
						),
					]
				),
			// create sequence of tuples, each combining an easing function with
			// a 2D position defining the bottom-left corner of a screen rect
			// (i.e. our grid layout), e.g. `[easeInCirc, [0.25, 0.5]]`...
			zip(EASINGS, normRange2d(COLS, ROWS, false, false))
		),
		// if fragment is outside any of the regions, return background color
		// (configured via shader uniforms, further below...)
		ret(unis.bgColor),
	];
};

// higher order function defining a shader function (also using
// thi.ng/shader-ast) to create a the plot of a single easing function within a
// given screen rect. the techniques used here give a glimpse of the full
// expressiveness offered by functional composition and compile-time code
// generation, in contrast to naive string concatenation/interpolation
// approaches. do check the console output to analyze the resulting GLSL!
const easingVisualization = (
	// the (shader) function to visualize
	easingFn: TaggedFn1<"float", "float">,
	// position & size of the view rect (UV coordinates)
	[minx, miny]: number[],
	[w, h]: number[],
	// shader uniforms
	unis: DemoUniforms
) =>
	defn(V4, null, [V2, V2, V2, F], (p, frag, mpos, t) => {
		let q: Vec2Sym, col: Vec4Sym;
		let d: FloatSym, d2: FloatSym;
		// pre-compute min/max bounds
		const bmin = vec2(minx + w * MARGIN, miny + h * MARGIN);
		const bmax = vec2(minx + w * (1 - MARGIN), miny + h * (1 - MARGIN));
		// pre-configure the plotting function
		const sampler = std.functionSampler(
			easingFn,
			std.functionDomainMapper(bmin, bmax)
		);
		// now the actual function body
		return [
			// compute the current (time based) curve position (in given screen space)
			(q = sym(madd(vec2(t, easingFn(t)), sub(bmax, bmin), bmin))),
			// compute distance to the dot at current position `q` (then threshold)
			(d = sym(step(float(unis.radius), distance(p, q)))),
			// compute distance to the curve/polyline (then threshold),
			(d2 = sym(
				smoothstep(
					unis.thickness,
					float(0),
					sampler(frag, unis.resolution)
				)
			)),
			// choose color for whichever shape is closest
			(col = sym(
				ternary(
					lte(d, d2),
					unis.dotColor,
					// compute gradient for curve color
					mix(
						unis.curveColor1,
						unis.curveColor2,
						$y(std.fitNorm(p, bmin, bmax))
					)
				)
			)),
			// if either `d` or `d2` are 1 return background color
			// (also choose bg color based on mouse position / hover state)
			ret(
				mix(
					col,
					// choose bg color based on mouse position
					mix(
						unis.bgColor,
						unis.bgHoverColor,
						float(
							std.isPointInRect(
								mpos,
								vec2(minx, miny),
								vec2(w, h)
							)
						)
					),
					min(d, d2)
				)
			),
		];
	});

// for simplicity, use a square aspect ratio for the canvas (use smallest side)
const W = Math.min(window.innerWidth, window.innerHeight - 32);

// create WebGL2 canvas
const { canvas, gl } = glCanvas({
	version: 2,
	width: W,
	height: W,
	parent: document.getElementById("app")!,
});

// init shadertoy scaffolding with canvas & main shader function
const toy = shaderToy({
	canvas,
	gl,
	main,
	// custom uniforms with default values
	uniforms: {
		thickness: [F, 0.2],
		radius: [F, 0.005],
		bgColor: [V4, [0.2, 0.2, 0.25, 1]],
		bgHoverColor: [V4, [0.3, 0.3, 0.35, 1]],
		curveColor1: [V4, [1, 1, 1, 1]],
		curveColor2: [V4, [0, 0.8, 1, 1]],
		dotColor: [V4, [1, 0, 0.8, 1]],
	},
	// GLSL code generation options (here to configure logger & float precision)
	// (see console for generated shader source code)
	opts: { logger: new ConsoleLogger("shader"), prec: 4 },
});

// kick off animation
toy.start();
